import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm  # tqdm是显示循环进度条的库

from env import CliffWalkingEnv, PrintPolicy


class Sarsa:
    """Sarsa算法"""

    def __init__(self, ncol, nrow, epsilon, alpha, gamma, n_action=4):
        self.Q_table = np.zeros([nrow * ncol, n_action])  # 初始化Q(s,a)表格
        self.n_action = n_action  # 动作个数
        self.alpha = alpha  # 学习率
        self.gamma = gamma  # 折扣因子
        self.epsilon = epsilon  # epsilon-贪婪策略中的参数

    def TakeAction(self, state):
        """根据 epsilon-贪心，选取下一步的动作"""
        if np.random.rand() < self.epsilon:
            action = np.random.randint(self.n_action)  # 随机选择动作
        else:
            action = np.argmax(self.Q_table[state])
        return action

    def PrintAction(self, state):
        Q_max = np.max(self.Q_table[state])
        a = [0 for _ in range(self.n_action)]
        for i in range(self.n_action):  # 若两个动作的价值一样,都会记录下来
            if self.Q_table[state, i] == Q_max:
                a[i] = 1
        return a

    def Update(self, s0, a0, r, s1, a1):
        td_error = r + self.gamma * self.Q_table[s1, a1] - self.Q_table[s0, a0]
        self.Q_table[s0, a0] += self.alpha * td_error


def PrintPolicy(agent, env, action_meaning, disaster=[], end=[]):
    for i in range(env.nrow):
        for j in range(env.ncol):
            if (i * env.ncol + j) in disaster:
                print("****", end=" ")
            elif (i * env.ncol + j) in end:
                print("EEEE", end=" ")
            else:
                a = agent.PrintAction(i * env.ncol + j)
                pi_str = ""
                for k in range(len(action_meaning)):
                    pi_str += action_meaning[k] if a[k] > 0 else "o"
                print(pi_str, end=" ")
        print()


if __name__ == "__main__":
    ncol = 12
    nrow = 4
    env = CliffWalkingEnv(ncol, nrow)

    np.random.seed(0)
    epsilon = 0.1
    alpha = 0.1
    gamma = 0.9
    num_episodes = 500  # 智能体在环境中运行的序列的数量

    agent = Sarsa(ncol, nrow, epsilon, alpha, gamma)

    num_episodes = 500  # 智能体在环境中运行的序列的数量

    return_list = []  # 记录每一条序列的回报
    for i in range(10):  # 显示10个进度条
        # tqdm的进度条功能
        with tqdm(total=int(num_episodes / 10), desc="Iteration %d" % i) as pbar:
            for i_episode in range(int(num_episodes / 10)):  # 每个进度条的序列数
                episode_return = 0
                state = env.Reset()
                action = agent.TakeAction(state)
                done = False
                while not done:
                    next_state, reward, done = env.Step(action)
                    next_action = agent.TakeAction(next_state)
                    episode_return += reward  # 这里回报的计算不进行折扣因子衰减
                    agent.Update(state, action, reward, next_state, next_action)
                    state = next_state
                    action = next_action
                return_list.append(episode_return)
                if (i_episode + 1) % 10 == 0:  # 每10条序列打印一下这10条序列的平均回报
                    pbar.set_postfix(
                        {
                            "episode": "%d" % (num_episodes / 10 * i + i_episode + 1),
                            "return": "%.3f" % np.mean(return_list[-10:]),
                        }
                    )
                pbar.update(1)

    episodes_list = list(range(len(return_list)))
    plt.plot(episodes_list, return_list)
    plt.xlabel("Episodes")
    plt.ylabel("Returns")
    plt.title("Sarsa on {}".format("Cliff Walking"))
    plt.show()

    action_meaning = ["^", "v", "<", ">"]
    print("Sarsa算法最终收敛得到的策略为：")
    PrintPolicy(agent, env, action_meaning, list(range(37, 47)), [47])
